#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 
constexpr uint32_t BLOCK_SIZE = 32;
constexpr uint32_t BLOCK_PER_UNIT = 8;
class KernelBallQuery{
public:
    __aicore__ inline KernelBallQuery() {}
    __aicore__ inline void Init(GM_ADDR xyz, GM_ADDR center_xyz, 
                                GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, 
                                GM_ADDR idx,
                                uint32_t B, uint32_t N, uint32_t M, 
                                uint32_t NAlign, uint32_t MAlign,
                                uint32_t baseB,
                                float min_radius, float max_radius, 
                                int32_t sample_num, int32_t sample_numAlign
                                )
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->B = B;
        this->N = N;
        this->M = M;
        this->NAlign = NAlign;
        this->MAlign = MAlign;
        this->baseB = baseB;
        this->sample_num = sample_num;
        this->sample_numAlign = sample_numAlign;
        this->neg_min = min_radius;
        this->neg_max = max_radius;
        xyzGm.SetGlobalBuffer((__gm__ float*)xyz);
        centerGm.SetGlobalBuffer((__gm__ float*)center_xyz);
        idxGm.SetGlobalBuffer((__gm__ int32_t*)idx);
        pipe.InitBuffer(inQueueC, BUFFER_NUM, (this->baseB) * (this->N) * 8 * sizeof(float));
        pipe.InitBuffer(inQueueXYZ, BUFFER_NUM, (this->baseB) * (this->M) * 8 * sizeof(float));
        pipe.InitBuffer(outQueueIDX, BUFFER_NUM,  (this->baseB) * (this->N) * (this->sample_numAlign) * sizeof(int32_t)); 
        this->maskLength = ((this->MAlign / (256/sizeof(float)))  + 1);
        pipe.InitBuffer(tmpBuffMASK, 3 * 32 * (this->maskLength)); 
        pipe.InitBuffer(workBuf, (256) * sizeof(float));
        this->workLocal = workBuf.Get<float>();
        pipe.InitBuffer(mdistBuf, 3 * (this->MAlign/8 + 1) * 64 * sizeof(float));  
        this->mdistLocal = mdistBuf.Get<float>();
        this->resulttLocal = mdistLocal[(this->MAlign/8 + 1) * 64];
        this->compresultLocal = resulttLocal[(this->MAlign/8 + 1) * 64];
        pipe.InitBuffer(tempBuf, 1 * (this->MAlign/64 + 1)* 64 * sizeof(float));
        LocalTensor<float> tempLocal = tempBuf.Get<float>();
        this->onesLocal = tempLocal;
    }
    __aicore__ inline void Process()
    {
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();    
        for(int inloop = block_id * (this->baseB); inloop <(this->B); inloop += core_num * (this->baseB)) 
        {
                this->calc_B = (this->baseB);
                this->xyz_offset = inloop * (this->M) * 3;
                this->c_offset = inloop * (this->N) * 3;
                this->idx_offset = inloop * (this->N) * (this->sample_num);
                if(inloop + baseB > (this->B)) 
                {
                    this->calc_B = ((this->B) - inloop);
                }
            this->xyz_blkcnt = (this->calc_B) * (this->M);
            this->c_blkcnt = (this->calc_B) * (this->N);
            this->idx_blkcnt = (this->calc_B) * (this->N); 
            this->idx_copyBytes = (this->sample_num) * sizeof(int32_t);
            CopyIn();
            Compute();
            CopyOut();            
        }
    }
private:
    __aicore__ inline void CopyIn()
    {
       LocalTensor<float> cLocal = inQueueC.AllocTensor<float>();
       LocalTensor<float> xyzLocal = inQueueXYZ.AllocTensor<float>();
        DataCopyExtParams xyz_copyParams{(this->xyz_blkcnt), 12, 0, 0, 0}; 
        DataCopyPadExtParams<float> padParams{true, 0, 0, 0}; 
        DataCopyPad(xyzLocal, xyzGm[this->xyz_offset], xyz_copyParams, padParams);         
        DataCopyExtParams c_copyParams{(this->c_blkcnt), 12, 0, 0, 0};
        DataCopyPad(cLocal, centerGm[this->c_offset], c_copyParams, padParams); 
       inQueueC.EnQue(cLocal);
       inQueueXYZ.EnQue(xyzLocal);
    }
    __aicore__ inline void Compute( )
    {
        LocalTensor<float> cLocal = inQueueC.DeQue<float>();
        LocalTensor<float> xyzLocal = inQueueXYZ.DeQue<float>();
        LocalTensor<int32_t> idxLocal = outQueueIDX.AllocTensor<int32_t>();  
        LocalTensor<uint16_t> maskLocal = tmpBuffMASK.Get<uint16_t>();
        LocalTensor<uint16_t> eq0_xmaskLocal = maskLocal[16 * maskLength];
        LocalTensor<uint16_t> geMin_xmaskLocal = maskLocal[32 * maskLength];  
        LocalTensor<uint16_t> ltMax_xmaskLocal = maskLocal[48 * maskLength];        
        for(int32_t b=0; b<(this->calc_B); b++)
        {
            int32_t loopcnt = (this->N);
            uint32_t iner_idx_cur = b * (this->N) * (this->sample_numAlign);
            uint32_t iner_c_bias = b * (this->N) * 8;
            uint32_t iner_xyz_bias = b * (this->M) * 8;
            for(int32_t i=0; i<loopcnt; i++){
                uint32_t idx_cur =iner_idx_cur + (i * (this->sample_numAlign));
                Duplicate(mdistLocal, (float)(0.0), ((this->M) * 8));
                Duplicate(resulttLocal, (float)(0.0), ((this->M) * 8));
                uint64_t mask[2] = { 506381209866536711, 0 }; 
                uint32_t repeat = (this->M) / 8;
                if((this->M % 8) != 0)
                {
                    repeat = repeat + 1;
                }
                Add(mdistLocal, mdistLocal, cLocal[iner_c_bias + i*8], mask, repeat, {1, 1, 0, 8 ,8 ,0});
                Add(resulttLocal, resulttLocal, xyzLocal[iner_xyz_bias], mask, repeat, {1, 1, 1, 8 ,8 ,8});
                Sub(mdistLocal, mdistLocal, resulttLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8}); 
                Mul(mdistLocal, mdistLocal, mdistLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8});
                BlockReduceSum<float>(mdistLocal, mdistLocal,  repeat, (this->M * 8),1, 1, 8); 
                CompareScalar(eq0_xmaskLocal, mdistLocal, (float)0.0, CMPMODE::EQ, ((this->M)));
                CompareScalar(geMin_xmaskLocal, mdistLocal, (float)(this->neg_min), CMPMODE::GE, ((this->M)));
                CompareScalar(ltMax_xmaskLocal, mdistLocal, (float)(this->neg_max), CMPMODE::LT, ((this->M)));
                And(geMin_xmaskLocal, geMin_xmaskLocal, ltMax_xmaskLocal, (16 * maskLength));
                Or(eq0_xmaskLocal, geMin_xmaskLocal, eq0_xmaskLocal, (16 * maskLength));
                Duplicate(this->onesLocal, (float)(1.0), (this->M));
                Select(resulttLocal, eq0_xmaskLocal, onesLocal, (float)(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, 64, repeat, { 1, 1, 1, 8, 0, 0});
                LocalTensor<int32_t>interpreLocal  = compresultLocal.ReinterpretCast<int32_t>();
                Cast(interpreLocal, resulttLocal, RoundMode::CAST_CEIL, (this->M));
                uint32_t cnt = 0;
                for (int j=0; j<(this->M); j++) {
                    if(interpreLocal.GetValue(j) == 1)
                    {
                        idxLocal.SetValue((idx_cur +cnt), (j));
                        if(cnt == 0)
                        {
                            Duplicate(idxLocal[idx_cur], j, (this->sample_num));
                        }
                        cnt = cnt + 1;
                    }
                    if(cnt == (this->sample_num))
                    {
                        break;
                    }
                }
            }
        }
        outQueueIDX.EnQue(idxLocal);
        inQueueC.FreeTensor(cLocal);
        inQueueXYZ.FreeTensor(xyzLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<int32_t> idxLocal = outQueueIDX.DeQue<int32_t>();
        DataCopyExtParams copyParams{(this->idx_blkcnt), (this->idx_copyBytes), 0, 0, 0}; 
        DataCopyPad(idxGm[this->idx_offset], idxLocal, copyParams);
        outQueueIDX.FreeTensor(idxLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueC, inQueueXYZ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueIDX;
    TBuf<TPosition::VECCALC>  tempBuf, workBuf, mdistBuf, tmpBuffMASK;
    GlobalTensor<float> xyzGm;
    GlobalTensor<float> centerGm;
    GlobalTensor<int32_t> idxGm;
    LocalTensor<float> workLocal;
    LocalTensor<float> onesLocal;
    LocalTensor<float> zeroLocal;
    LocalTensor<float> mdistLocal;
    LocalTensor<float> resulttLocal;
    LocalTensor<float> compresultLocal;
    uint32_t xyz_offset;
    uint16_t xyz_blkcnt;
    uint16_t c_blkcnt;
    uint32_t c_offset;
    uint16_t idx_blkcnt;
    uint32_t idx_offset;
    uint32_t idx_copyBytes;
    uint32_t calc_B;
    uint32_t maskLength;
    uint32_t B;
    uint32_t N;
    uint32_t M;
    uint32_t NAlign;
    uint32_t MAlign;
    uint32_t baseB;
    float neg_min; 
    float neg_max; 
    int32_t sample_num;
    int32_t sample_numAlign;
    uint32_t tileNum;
    uint32_t tileLength;
};
class KernelBallQuery_fp16{
public:
    __aicore__ inline KernelBallQuery_fp16() {}
    __aicore__ inline void Init(GM_ADDR xyz, GM_ADDR center_xyz, 
                                GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, 
                                GM_ADDR idx,
                                uint32_t B, uint32_t N, uint32_t M, 
                                uint32_t NAlign, uint32_t MAlign,
                                uint32_t baseB,
                                float min_radius, float max_radius, 
                                int32_t sample_num, int32_t sample_numAlign
                                )
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->B = B;
        this->N = N;
        this->M = M;
        this->NAlign = NAlign;
        this->MAlign = MAlign;
        this->baseB = baseB;
        this->sample_num = sample_num;
        this->sample_numAlign = sample_numAlign;
        this->neg_min = min_radius;
        this->neg_max = max_radius;
        xyzGm.SetGlobalBuffer((__gm__ half*)xyz);
        centerGm.SetGlobalBuffer((__gm__ half*)center_xyz);
        idxGm.SetGlobalBuffer((__gm__ int32_t*)idx);
        pipe.InitBuffer(inQueueC, BUFFER_NUM, (this->baseB) * ((this->N)* 16 * sizeof(half)));
        pipe.InitBuffer(inQueueXYZ, BUFFER_NUM, (this->baseB) * ((this->M) * 16 * sizeof(half)));
        pipe.InitBuffer(outQueueIDX, BUFFER_NUM, (this->baseB) * (this->N) * (this->sample_numAlign) * sizeof(int32_t)); 
        this->maskLength = ((this->MAlign / (256/sizeof(half)))  + 1);
        pipe.InitBuffer(tmpBuffMASK, 3 * 256 * (this->maskLength) * sizeof(uint8_t)); 
        pipe.InitBuffer(workBuf, (256) * sizeof(half));
        this->workLocal = workBuf.Get<half>();
        pipe.InitBuffer(mdistBuf, 3 * (this->MAlign/8 + 1) * 128 * sizeof(half));  
        this->mdistLocal = mdistBuf.Get<half>();
        this->resulttLocal = mdistLocal[(this->MAlign/8 + 1) * 128];
        this->compresultLocal = resulttLocal[(this->MAlign/8 + 1) * 128]; 
        pipe.InitBuffer(tempBuf, 1 * (this->MAlign / 128 + 1)* 256);
        LocalTensor<half> tempLocal = tempBuf.Get<half>();
        this->onesLocal = tempLocal;
        pipe.InitBuffer(tempCompare, 2*128*sizeof(half));
    }
    __aicore__ inline void Process()
    {
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();    
        for(int inloop = block_id * (this->baseB); inloop <(this->B); inloop += core_num * (this->baseB)) 
        {
                this->calc_B = (this->baseB);
                this->xyz_offset = inloop * (this->M) * 3;
                this->c_offset = inloop * (this->N) * 3;
                this->idx_offset = inloop * (this->N) * (this->sample_num);
                if(inloop + baseB > (this->B)) 
                {
                    this->calc_B = ((this->B) - inloop);
                }
            this->xyz_blkcnt = (this->calc_B) * (this->M);
            this->c_blkcnt = (this->calc_B) * (this->N);
            this->idx_blkcnt = (this->calc_B) * (this->N); 
            this->idx_copyBytes = (this->sample_num) * sizeof(int32_t);
            CopyIn();
            Compute();
            CopyOut();            
        }
    }
private:
    __aicore__ inline void CopyIn()
    {
       LocalTensor<half> cLocal = inQueueC.AllocTensor<half>();
       LocalTensor<half> xyzLocal = inQueueXYZ.AllocTensor<half>();
        DataCopyExtParams xyz_copyParams{(this->xyz_blkcnt), 6, 0, 0, 0}; 
        DataCopyPadExtParams<half> padParams{true, 0, 0, 0}; 
        DataCopyPad(xyzLocal, xyzGm[this->xyz_offset], xyz_copyParams, padParams);         
        DataCopyExtParams c_copyParams{(this->c_blkcnt), 6, 0, 0, 0};
        DataCopyPad(cLocal, centerGm[this->c_offset], c_copyParams, padParams); 
       inQueueC.EnQue(cLocal);
       inQueueXYZ.EnQue(xyzLocal);
    }
    __aicore__ inline void Compute( )
    {
        LocalTensor<half> cLocal = inQueueC.DeQue<half>();
        LocalTensor<half> xyzLocal = inQueueXYZ.DeQue<half>();
        LocalTensor<int32_t> idxLocal = outQueueIDX.AllocTensor<int32_t>();  
        LocalTensor<uint8_t> maskLocal = tmpBuffMASK.Get<uint8_t>();
        LocalTensor<uint8_t> eq0_xmaskLocal = maskLocal;
        LocalTensor<uint8_t> geMin_xmaskLocal = maskLocal[256];  
        LocalTensor<uint8_t> ltMax_xmaskLocal = maskLocal[512];        
        for(int32_t b=0; b<(this->calc_B); b++)
        {
            int32_t loopcnt = (this->N);
            uint32_t iner_idx_cur = b * (this->N) * (this->sample_numAlign);
            uint32_t iner_c_bias = b * (this->N) * 16;
            uint32_t iner_xyz_bias = b * (this->M) * 16;
            for(int32_t i=0; i<loopcnt; i++){
                uint32_t idx_cur =iner_idx_cur + (i * (this->sample_numAlign));
                Duplicate(mdistLocal, (half)(0.0), ((this->M) * 16));
                Duplicate(resulttLocal, (half)(0.0), ((this->M) * 16));
                uint64_t mask[2] = { 1970354902204423, 1970354902204423 }; 
                uint32_t repeat = (this->M) / 8; 
                if((this->M % 8) != 0)
                {
                    repeat = repeat + 1;
                }
                Add(mdistLocal, mdistLocal, cLocal[iner_c_bias + i*16], mask, repeat, {1, 1, 0, 8 ,8 ,0});
                Add(resulttLocal, resulttLocal, xyzLocal[iner_xyz_bias], mask, repeat, {1, 1, 1, 8 ,8 ,8});
                Sub(mdistLocal, mdistLocal, resulttLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8}); 
                Mul(mdistLocal, mdistLocal, mdistLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8});
                BlockReduceSum<half>(mdistLocal, mdistLocal,  repeat, (this->M * (BLOCK_SIZE/sizeof(half))),1, 1, 8); 
                DumpTensor(mdistLocal, 1008, 32);
                LocalTensor<half> const1_Local =  tempCompare.Get<half>();
                LocalTensor<half> const2_Local = const1_Local[128];
                Duplicate(const1_Local, (half)(0.0), ((128)));
                Compare(eq0_xmaskLocal, mdistLocal, const1_Local, CMPMODE::EQ, ((128)));
                Duplicate(const1_Local, (half)(100.0), ((128)));
                Compare(geMin_xmaskLocal, mdistLocal, const1_Local, CMPMODE::GE, ((128)));
                Duplicate(const1_Local, (half)(2500.0), ((128)));
                Compare(ltMax_xmaskLocal, mdistLocal, const1_Local, CMPMODE::LT, ((128)));
                And(geMin_xmaskLocal, geMin_xmaskLocal, ltMax_xmaskLocal, (16 * maskLength));
                Or(eq0_xmaskLocal, geMin_xmaskLocal, eq0_xmaskLocal, (16 * maskLength));
                Duplicate(this->onesLocal, (half)(1.0), (128));
                uint32_t selcet_repeat = ((this->M) / 128 + 1);
                Select(resulttLocal, eq0_xmaskLocal, onesLocal, (half)(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, 128, selcet_repeat, { 1, 1, 1, 8, 0, 0});
                LocalTensor<int32_t>interpreLocal  = compresultLocal.ReinterpretCast<int32_t>();
                Cast(interpreLocal, resulttLocal, RoundMode::CAST_CEIL, (this->M));
                DumpTensor(resulttLocal, 2007, 32);
                uint32_t cnt = 0;
                for (int j=0; j<(this->M); j++) {
                    if(interpreLocal.GetValue(j) == 1)
                    {
                        idxLocal.SetValue((idx_cur +cnt), (j));
                        if(cnt == 0)
                        {
                            Duplicate(idxLocal[idx_cur], j, (this->sample_num));
                        }
                        cnt = cnt + 1;
                    }
                    if(cnt == (this->sample_num))
                    {
                        break;
                    }
                }
            }
        }
        outQueueIDX.EnQue(idxLocal);
        inQueueC.FreeTensor(cLocal);
        inQueueXYZ.FreeTensor(xyzLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<int32_t> idxLocal = outQueueIDX.DeQue<int32_t>();
        DataCopyExtParams copyParams{(this->idx_blkcnt), (this->idx_copyBytes), 0, 0, 0}; 
        DataCopyPad(idxGm[this->idx_offset], idxLocal, copyParams);
        outQueueIDX.FreeTensor(idxLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueC, inQueueXYZ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueIDX;
    TBuf<TPosition::VECCALC>  tempBuf, workBuf, mdistBuf, tmpBuffMASK, tempCompare;
    GlobalTensor<half> xyzGm;
    GlobalTensor<half> centerGm;
    GlobalTensor<int32_t> idxGm;
    LocalTensor<half> workLocal;
    LocalTensor<half> onesLocal;
    LocalTensor<half> zeroLocal;
    LocalTensor<half> mdistLocal;
    LocalTensor<half> resulttLocal;
    LocalTensor<half> compresultLocal;
    uint32_t xyz_offset;
    uint16_t xyz_blkcnt;
    uint16_t c_blkcnt;
    uint32_t c_offset;
    uint16_t idx_blkcnt;
    uint32_t idx_offset;
    uint32_t idx_copyBytes;
    uint32_t calc_B;
    uint32_t maskLength;
    uint32_t B;
    uint32_t N;
    uint32_t M;
    uint32_t NAlign;
    uint32_t MAlign;
    uint32_t baseB;
    float neg_min; 
    float neg_max; 
    int32_t sample_num;
    int32_t sample_numAlign;
    uint32_t tileNum;
    uint32_t tileLength;
};
/*
class KernelBallQuery_fp16{
public:
    __aicore__ inline KernelBallQuery_fp16() {}
    __aicore__ inline void Init(GM_ADDR xyz, GM_ADDR center_xyz, 
                                GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, 
                                GM_ADDR idx,
                                uint32_t B, uint32_t N, uint32_t M, 
                                uint32_t NAlign, uint32_t MAlign,
                                uint32_t baseB,
                                float min_radius, float max_radius, 
                                int32_t sample_num, int32_t sample_numAlign
                                )
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->B = B;
        this->N = N;
        this->M = M;
        this->NAlign = NAlign;
        this->MAlign = MAlign;
        this->baseB = baseB;
        this->sample_num = sample_num;
        this->sample_numAlign = sample_numAlign;
        this->neg_min = min_radius;
        this->neg_max = max_radius;
        xyzGm.SetGlobalBuffer((__gm__ half*)xyz);
        centerGm.SetGlobalBuffer((__gm__ half*)center_xyz);
        idxGm.SetGlobalBuffer((__gm__ int32_t*)idx);
        pipe.InitBuffer(inQueueC, BUFFER_NUM, (this->baseB) * (this->N) * 8 * sizeof(half));
        pipe.InitBuffer(inQueueXYZ, BUFFER_NUM, (this->baseB) * (this->M) * 8 * sizeof(half));
        pipe.InitBuffer(tempBuffC, (this->baseB) * (this->N) * 8 * sizeof(float));
        pipe.InitBuffer(tempBuffXYZ, (this->baseB) * (this->M) * 8 * sizeof(float));        
        pipe.InitBuffer(outQueueIDX, BUFFER_NUM, (this->N) * (this->sample_numAlign) * sizeof(int32_t)); 
        this->maskLength = ((this->MAlign / (256/sizeof(float)))  + 1);
        pipe.InitBuffer(tmpBuffMASK, 3 * 32 * (this->maskLength)); 
        pipe.InitBuffer(workBuf, (256) * sizeof(float));
        this->workLocal = workBuf.Get<float>();
        pipe.InitBuffer(mdistBuf, 3 * (this->MAlign/8 + 1) * 64 * sizeof(float));  
        this->mdistLocal = mdistBuf.Get<float>();
        this->resulttLocal = mdistLocal[(this->MAlign/8 + 1) * 64];
        this->compresultLocal = resulttLocal[(this->MAlign/8 + 1) * 64];
        pipe.InitBuffer(tempBuf, 1 * (this->MAlign+64)* sizeof(float));
        LocalTensor<float> tempLocal = tempBuf.Get<float>();
        this->onesLocal = tempLocal;
    }
    __aicore__ inline void Process()
    {
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();    
        for(int inloop = block_id * (this->baseB); inloop <(this->B); inloop += core_num * (this->baseB)) 
        {
                this->calc_B = (this->baseB);
                this->xyz_offset = inloop * (this->M) * 3;
                this->c_offset = inloop * (this->N) * 3;
                this->idx_offset = inloop * (this->N) * (this->sample_num);
                if(inloop + baseB > (this->B)) 
                {
                    this->calc_B = ((this->B) - inloop);
                }
            this->xyz_blkcnt = (this->calc_B) * (this->M);
            this->c_blkcnt = (this->calc_B) * (this->N);
            this->idx_blkcnt = (this->calc_B) * (this->N); 
            this->idx_copyBytes = (this->sample_num) * sizeof(float);
            CopyIn();
            Compute();
            CopyOut();            
        }
    }
private:
    __aicore__ inline void CopyIn()
    {
       LocalTensor<half> cLocal = inQueueC.AllocTensor<half>();
       LocalTensor<half> xyzLocal = inQueueXYZ.AllocTensor<half>();
        DataCopyExtParams xyz_copyParams{(this->xyz_blkcnt), 6, 0, 0, 0}; 
        DataCopyPadExtParams<half> padParams{true, 0, 0, 0}; 
        DataCopyPad(xyzLocal, xyzGm[this->xyz_offset], xyz_copyParams, padParams);         
        DataCopyExtParams c_copyParams{(this->c_blkcnt), 6, 0, 0, 0};
        DataCopyPad(cLocal, centerGm[this->c_offset], c_copyParams, padParams); 
       inQueueC.EnQue(cLocal);
       inQueueXYZ.EnQue(xyzLocal);
    }
    __aicore__ inline void Compute( )
    {
        LocalTensor<half> in_cLocal = inQueueC.DeQue<half>();
        LocalTensor<half> in_xyzLocal = inQueueXYZ.DeQue<half>();
        LocalTensor<int32_t> idxLocal = outQueueIDX.AllocTensor<int32_t>();  
        LocalTensor<uint16_t> maskLocal = tmpBuffMASK.Get<uint16_t>();
        LocalTensor<uint16_t> eq0_xmaskLocal = maskLocal[16 * maskLength];
        LocalTensor<uint16_t> geMin_xmaskLocal = maskLocal[32 * maskLength];  
        LocalTensor<uint16_t> ltMax_xmaskLocal = maskLocal[48 * maskLength];  
        LocalTensor<float> cLocal = tempBuffC.Get<float>();
        LocalTensor<float> xyzLocal = tempBuffXYZ.Get<float>();
        Cast(cLocal, in_cLocal, RoundMode::CAST_CEIL, );
        for(int32_t b=0; b<(this->calc_B); b++)
        {
            int32_t loopcnt = (this->N);
            uint32_t iner_idx_cur = b * (this->N) * (this->sample_numAlign);
            uint32_t iner_c_bias = b * (this->N) * 8;
            uint32_t iner_xyz_bias = b * (this->M) * 8;
            for(int32_t i=0; i<loopcnt; i++){
                uint32_t idx_cur =iner_idx_cur + (i * (this->sample_numAlign));
                Duplicate(mdistLocal, (float)(0.0), ((this->M) * 8));
                Duplicate(resulttLocal, (float)(0.0), ((this->M) * 8));
                uint64_t mask[2] = { 506381209866536711, 0 }; 
                uint32_t repeat = (this->M) / 8;
                if((this->M % 8) != 0)
                {
                    repeat = repeat + 1;
                }
                Add(mdistLocal, mdistLocal, cLocal[iner_c_bias + i*8], mask, repeat, {1, 1, 0, 8 ,8 ,0});
                Add(resulttLocal, resulttLocal, xyzLocal[iner_xyz_bias], mask, repeat, {1, 1, 1, 8 ,8 ,8});
                Sub(mdistLocal, mdistLocal, resulttLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8}); 
                Mul(mdistLocal, mdistLocal, mdistLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8});
                BlockReduceSum<float>(mdistLocal, mdistLocal,  repeat, (this->M * 8),1, 1, 8); 
                CompareScalar(eq0_xmaskLocal, mdistLocal, (float)0.0, CMPMODE::EQ, ((this->M)));
                CompareScalar(geMin_xmaskLocal, mdistLocal, (float)(this->neg_min), CMPMODE::GE, ((this->M)));
                CompareScalar(ltMax_xmaskLocal, mdistLocal, (float)(this->neg_max), CMPMODE::LT, ((this->M)));
                And(geMin_xmaskLocal, geMin_xmaskLocal, ltMax_xmaskLocal, (16 * maskLength));
                Or(eq0_xmaskLocal, geMin_xmaskLocal, eq0_xmaskLocal, (16 * maskLength));
                Duplicate(this->onesLocal, (float)(1.0), (this->M));
                Select(resulttLocal, eq0_xmaskLocal, onesLocal, (float)(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, 64, repeat, { 1, 1, 1, 8, 0, 0});
                LocalTensor<int32_t>interpreLocal  = compresultLocal.ReinterpretCast<int32_t>();
                Cast(interpreLocal, resulttLocal, RoundMode::CAST_CEIL, (this->M));
                uint32_t cnt = 0;
                for (int j=0; j<(this->M); j++) {
                    if(interpreLocal.GetValue(j) == 1)
                    {
                        idxLocal.SetValue((idx_cur +cnt), (j));
                        if(cnt == 0)
                        {
                            Duplicate(idxLocal[idx_cur], j, (this->sample_num));
                        }
                        cnt = cnt + 1;
                    }
                    if(cnt == (this->sample_num))
                    {
                        break;
                    }
                }
            }
        }
        outQueueIDX.EnQue(idxLocal);
        inQueueC.FreeTensor(cLocal);
        inQueueXYZ.FreeTensor(xyzLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<int32_t> idxLocal = outQueueIDX.DeQue<int32_t>();
        DataCopyExtParams copyParams{(this->idx_blkcnt), (this->idx_copyBytes), 0, 0, 0}; 
        DataCopyPad(idxGm[this->idx_offset], idxLocal, copyParams);
        outQueueIDX.FreeTensor(idxLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueC, inQueueXYZ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueIDX;
    TBuf<TPosition::VECCALC>  tempBuf, workBuf, mdistBuf, tmpBuffMASK, tempBuffC, tempBuffXYZ;
    GlobalTensor<float> xyzGm;
    GlobalTensor<float> centerGm;
    GlobalTensor<int32_t> idxGm;
    LocalTensor<float> workLocal;
    LocalTensor<float> onesLocal;
    LocalTensor<float> zeroLocal;
    LocalTensor<float> mdistLocal;
    LocalTensor<float> resulttLocal;
    LocalTensor<float> compresultLocal;
    uint32_t xyz_offset;
    uint16_t xyz_blkcnt;
    uint16_t c_blkcnt;
    uint32_t c_offset;
    uint16_t idx_blkcnt;
    uint32_t idx_offset;
    uint32_t idx_copyBytes;
    uint32_t calc_B;
    uint32_t maskLength;
    uint32_t B;
    uint32_t N;
    uint32_t M;
    uint32_t NAlign;
    uint32_t MAlign;
    uint32_t baseB;
    float neg_min; 
    float neg_max; 
    int32_t sample_num;
    int32_t sample_numAlign;
    uint32_t tileNum;
    uint32_t tileLength;
};
*/
class KernelBallQuery_b{
public:
    __aicore__ inline KernelBallQuery_b() {}
    __aicore__ inline void Init(GM_ADDR xyz, GM_ADDR center_xyz, 
                                GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, 
                                GM_ADDR idx,
                                uint32_t B, uint32_t N, uint32_t M, 
                                uint32_t NAlign, uint32_t MAlign,
                                uint32_t baseB,
                                float min_radius, float max_radius, 
                                int32_t sample_num, int32_t sample_numAlign,
                                uint32_t batch_shape
                                )
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->B = B;
        this->N = N;
        this->M = M;
        this->NAlign = NAlign;
        this->MAlign = MAlign;
        this->baseB = baseB;
        this->sample_num = sample_num;
        this->sample_numAlign = sample_numAlign;
        this->neg_min = min_radius;
        this->neg_max = max_radius;
        this->batch_shape = batch_shape;
        xyzGm.SetGlobalBuffer((__gm__ float*)xyz);
        centerGm.SetGlobalBuffer((__gm__ float*)center_xyz);
        xyz_batchGm.SetGlobalBuffer((__gm__ int32_t*)xyz_batch_cnt);
        c_batchGm.SetGlobalBuffer((__gm__ int32_t*)center_xyz_batch_cnt);
        idxGm.SetGlobalBuffer((__gm__ int32_t*)idx);
       pipe.InitBuffer(inQueueC, BUFFER_NUM, (this->baseB) * (this->N) * 8 * sizeof(float));
        pipe.InitBuffer(inQueueXYZ, BUFFER_NUM, (this->baseB) * (this->M) * 8 * sizeof(float));
        pipe.InitBuffer(outQueueIDX, BUFFER_NUM, (this->N) * (this->sample_numAlign) * sizeof(int32_t)); 
        this->maskLength = ((this->MAlign / (256/sizeof(float)))  + 1);
        pipe.InitBuffer(tmpBuffMASK, 3 * 32 * (this->maskLength)); 
        pipe.InitBuffer(workBuf, (256) * sizeof(float));
        this->workLocal = workBuf.Get<float>();
        pipe.InitBuffer(mdistBuf, 3 * (this->MAlign/8 + 1) * 64 * sizeof(float));  
        this->mdistLocal = mdistBuf.Get<float>();
        this->resulttLocal = mdistLocal[(this->MAlign/8 + 1) * 64];
        this->compresultLocal = resulttLocal[(this->MAlign/8 + 1) * 64];
        pipe.InitBuffer(tempBuf, 1 * (this->MAlign+64)* sizeof(float));
        LocalTensor<float> tempLocal = tempBuf.Get<float>();
        this->onesLocal = tempLocal;        
        this->xyz_batch_bias[0] = 0;
        this->c_batch_bias[0] = 0;
        for (int i=0; i<batch_shape-1; i++) {
            this->xyz_batchArr[i] = xyz_batchGm.GetValue(i);
            this->c_batchArr[i] = c_batchGm.GetValue(i);
            this->xyz_batch_bias[i+1] = xyz_batch_bias[i-1] + xyz_batchArr[i];
            this->c_batch_bias[i+1] = c_batch_bias[i-1] + c_batchArr[i];
        }
        this->xyz_batchArr[batch_shape-1] = xyz_batchGm.GetValue(batch_shape-1);
        this->c_batchArr[batch_shape-1] = c_batchGm.GetValue(batch_shape-1);
        this->xyz_blkcnt = (this->M);
        (this->c_blkcnt) = (this->N);
        this->idx_blkcnt = (this->N);
       LocalTensor<float> cLocal = inQueueC.AllocTensor<float>();
       LocalTensor<float> xyzLocal = inQueueXYZ.AllocTensor<float>();
        DataCopyExtParams xyz_copyParams{(this->xyz_blkcnt), 12, 0, 0, 0}; 
        DataCopyPadExtParams<float> padParams{true, 0, 0, 0}; 
        DataCopyPad(xyzLocal, xyzGm[this->xyz_offset], xyz_copyParams, padParams);         
        DataCopyExtParams c_copyParams{(this->c_blkcnt), 12, 0, 0, 0};
        DataCopyPad(cLocal, centerGm[this->c_offset], c_copyParams, padParams); 
       inQueueC.EnQue(cLocal);
       inQueueXYZ.EnQue(xyzLocal);        
        cLocal = inQueueC.DeQue<float>();
        xyzLocal = inQueueXYZ.DeQue<float>();
        LocalTensor<int32_t> idxLocal = outQueueIDX.AllocTensor<int32_t>();  
        LocalTensor<uint16_t> maskLocal = tmpBuffMASK.Get<uint16_t>();
        LocalTensor<uint16_t> eq0_xmaskLocal = maskLocal[16 * maskLength];
        LocalTensor<uint16_t> geMin_xmaskLocal = maskLocal[32 * maskLength];  
        LocalTensor<uint16_t> ltMax_xmaskLocal = maskLocal[48 * maskLength];      
        for(int32_t b=0; b<(this->batch_shape); b++)
        {
            int32_t loopcnt = c_batchArr[b];
            uint32_t iner_idx_cur = this->c_batch_bias[b] * (this->sample_numAlign);
            uint32_t iner_c_bias = this->c_batch_bias[b];
            uint32_t iner_xyz_bias = this->xyz_batch_bias[b];
            for(int32_t i=0; i<loopcnt; i++){
                uint32_t idx_cur =iner_idx_cur + (i * (this->sample_numAlign));
                Duplicate(mdistLocal, (float)(0.0), ((this->M) * 8));
                Duplicate(resulttLocal, (float)(0.0), ((this->M) * 8));
                uint64_t mask[2] = { 506381209866536711, 0 }; 
                uint32_t repeat = (xyz_batchArr[i]) / 8;
                if(((this->xyz_batchArr[i]) % 8) != 0)
                {
                    repeat = repeat + 1;
                }
                Add(mdistLocal, mdistLocal, cLocal[iner_c_bias + i*8], mask, repeat, {1, 1, 0, 8 ,8 ,0});
                Add(resulttLocal, resulttLocal, xyzLocal[iner_xyz_bias], mask, repeat, {1, 1, 1, 8 ,8 ,8});
                Sub(mdistLocal, mdistLocal, resulttLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8}); 
                Mul(mdistLocal, mdistLocal, mdistLocal, mask, repeat, {1, 1, 1, 8 ,8 ,8});
                BlockReduceSum<float>(mdistLocal, mdistLocal,  repeat, ((this->xyz_batchArr[i]) * 8),1, 1, 8); 
                CompareScalar(eq0_xmaskLocal, mdistLocal, (float)0.0, CMPMODE::EQ, ((this->xyz_batchArr[i])));
                CompareScalar(geMin_xmaskLocal, mdistLocal, (float)(this->neg_min), CMPMODE::GE, ((this->xyz_batchArr[i])));
                CompareScalar(ltMax_xmaskLocal, mdistLocal, (float)(this->neg_max), CMPMODE::LT, ((this->xyz_batchArr[i])));
                And(geMin_xmaskLocal, geMin_xmaskLocal, ltMax_xmaskLocal, (16 * maskLength));
                Or(eq0_xmaskLocal, geMin_xmaskLocal, eq0_xmaskLocal, (16 * maskLength));
                Duplicate(this->onesLocal, (float)(1.0), (this->xyz_batchArr[i]));
                Select(resulttLocal, eq0_xmaskLocal, onesLocal, (float)(0.0), SELMODE::VSEL_TENSOR_SCALAR_MODE, 64, repeat, { 1, 1, 1, 8, 0, 0});
                LocalTensor<int32_t>interpreLocal  = compresultLocal.ReinterpretCast<int32_t>();
                Cast(interpreLocal, resulttLocal, RoundMode::CAST_CEIL, (this->xyz_batchArr[i]));
                uint32_t cnt = 0;
                for (int j=0; j<((this->xyz_batchArr[i])); j++) {
                    if(interpreLocal.GetValue(j) == 1)
                    {
                        idxLocal.SetValue((idx_cur +cnt), (j));
                        if(cnt == 0)
                        {
                            Duplicate(idxLocal[idx_cur], j, (this->sample_num));
                        }
                        cnt = cnt + 1;
                    }
                    if(cnt == (this->sample_num))
                    {
                        break;
                    }
                }
            }
        }
        outQueueIDX.EnQue(idxLocal);
        inQueueC.FreeTensor(cLocal);
        inQueueXYZ.FreeTensor(xyzLocal);
        idxLocal = outQueueIDX.DeQue<int32_t>();
        DataCopyExtParams copyParams{(this->idx_blkcnt), (this->idx_copyBytes), 0, 0, 0}; 
        DataCopyPad(idxGm, idxLocal, copyParams);
        outQueueIDX.FreeTensor(idxLocal);        
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueC, inQueueXYZ;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueIDX;
    TBuf<TPosition::VECCALC>  tempBuf, workBuf, mdistBuf, tmpBuffMASK;
    GlobalTensor<float> xyzGm;
    GlobalTensor<float> centerGm;
    GlobalTensor<int32_t> xyz_batchGm;
    GlobalTensor<int32_t> c_batchGm;
    GlobalTensor<int32_t> idxGm;
    LocalTensor<float> workLocal;
    LocalTensor<float> onesLocal;
    LocalTensor<float> zeroLocal;
    LocalTensor<float> mdistLocal;
    LocalTensor<float> resulttLocal;
    LocalTensor<float> compresultLocal;
    int32_t xyz_batchArr[10];
    int32_t c_batchArr[10];
    int32_t xyz_batch_bias[10];
    int32_t c_batch_bias[10];    
    uint32_t xyz_offset;
    uint16_t xyz_blkcnt;
    uint16_t c_blkcnt;
    uint32_t c_offset;
    uint16_t idx_blkcnt;
    uint32_t idx_offset;
    uint32_t idx_copyBytes;
    uint32_t calc_B;
    uint32_t maskLength;
    uint32_t B;
    uint32_t N;
    uint32_t M;
    uint32_t NAlign;
    uint32_t MAlign;
    uint32_t baseB;
    float neg_min; 
    float neg_max; 
    int32_t sample_num;
    int32_t sample_numAlign;
    uint32_t batch_shape;
};
extern "C" __global__ __aicore__ void ball_query(GM_ADDR xyz, GM_ADDR center_xyz, GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, GM_ADDR idx, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if(TILING_KEY_IS(1)){
        KernelBallQuery OP;
        OP.Init(xyz, center_xyz,
               xyz_batch_cnt, center_xyz_batch_cnt,
               idx,
               tiling_data.B, tiling_data.N, tiling_data.M, 
               tiling_data.NAlign, tiling_data.MAlign,
               tiling_data.baseB,
               tiling_data.neg_min, tiling_data.neg_max,
               tiling_data.sample_num, tiling_data.sample_numAlign
               );
        OP.Process();
    }
    else if(TILING_KEY_IS(2)){
        KernelBallQuery_fp16 OP_16;
        OP_16.Init(xyz, center_xyz,
               xyz_batch_cnt, center_xyz_batch_cnt,
               idx,
               tiling_data.B, tiling_data.N, tiling_data.M, 
               tiling_data.NAlign, tiling_data.MAlign,
               tiling_data.baseB,
               tiling_data.neg_min, tiling_data.neg_max,
               tiling_data.sample_num, tiling_data.sample_numAlign
               );
        OP_16.Process();
    }  
    else if(TILING_KEY_IS(3)){
        KernelBallQuery_b OP_b;
        OP_b.Init(xyz, center_xyz,
               xyz_batch_cnt, center_xyz_batch_cnt,
               idx,
               tiling_data.B, tiling_data.N, tiling_data.M, 
               tiling_data.NAlign, tiling_data.MAlign,
               tiling_data.baseB,
               tiling_data.neg_min, tiling_data.neg_max,
               tiling_data.sample_num, tiling_data.sample_numAlign,
               tiling_data.batch_shape
               );
    }  
}
